#!/usr/bin/python

"""Copies a MySQL database including all relevant privileges from one host to
another.

Author: Dave Challis
Email: suicas@gmail.com
Version: 1.1.0
"""

import sys
import MySQLdb
import re
import getpass
import pexpect
import tempfile
import os
from optparse import OptionParser

class Settings(object):
    """Class used to hold global settings."""
    # paths to needed binaries
    mysql_path = '/usr/bin/mysql'
    mysqldump_path = '/usr/bin/mysqldump'

    # path to temporary directory - specify None to use system default
    temp_dir = None

    # extra arguments to pass to mysql/mysqldump, should be a list of strings
    mysql_extra_args = []
    mysqldump_extra_args = []

    # 0 = normal, -1 = quiet, 1 = verbose
    verbosity = 0

    # set to True to answer yes to each question when prompted
    force = False

    # set to True to print changes without running them on any DB
    dry_run = False


def vprint(string, min_verbosity):
    """Conditionally prints string passed to it based on verbosity level"""
    if min_verbosity <= Settings.verbosity:
        print string,


def get_option_parser():
    """Get a handler for command line arguments"""

    usage = """%prog [OPTIONS]
Copies a MySQL database and relevant privileges from one MySQL server to
another.

Run with '%prog -i' for an interactive session which prompts for all options as it runs.

Note: This isn't currently transactional, so cancelling halfway through a copy
may leave incomplete data on the destination server.  Run with --dry-run if
you need to test copying first.
"""
    parser = OptionParser(usage=usage)

    help = 'Interactively prompt for all information'
    parser.add_option('-i', '--interactive', help=help, action='store_true')

    help = 'Execute all SQL without prompting for confirmation'
    parser.add_option('-f', '--force', help=help, action='store_true')

    help = """Suppress all output except where input is needed.
To completely suppress all output, specify all needed options on the command
line, and use the --force flag.
    """
    parser.add_option('-q', '--quiet', help=help, action='store_true')
    
    help = 'Print more verbose output at each stage'
    parser.add_option('-v', '--verbose', help=help, action='store_true')
    
    help = 'Do not execute any SQL that modifies data on servers'
    parser.add_option('--dry-run', help=help, action='store_true')

    help = 'Source host to copy database and privileges from'
    parser.add_option('--source-host', help=help, metavar='HOST')

    help = 'User to connect to source host as'
    parser.add_option('--source-user', help=help, metavar='USER')

    help = """Password to connect to the source host with.
        WARNING: this is insecure as password will show up in the process list.
    """
    parser.add_option('--source-password', help=help, metavar='PASSWORD')

    help = 'Database to copy from the source host'
    parser.add_option('--source-database', help=help, metavar='DBNAME')

    help = 'Destination host to copy database and privileges to'
    parser.add_option('--destination-host', help=help, metavar='HOST')

    help = 'User to connect to destination host as'
    parser.add_option('--destination-user', help=help, metavar='USER')

    help = """Password to connect to the destination host with.
        WARNING: this is insecure as password will show up in the process list.
    """
    parser.add_option('--destination-password', help=help, metavar='PASSWORD')

    help = """Path to alternative temporary directory.
        Default is operating system temp dir (e.g. /tmp on linux).
    """
    parser.add_option('--temp-dir', help=help, metavar='PATH')

    help = """A command line option to pass to the mysql binary.
    Pass these unchanged, or quoted if you need to use a space,
e.g.:                                     
    --mysql-option=--connect_timeout=20                      
    --mysql-option=-Q                           
    --mysql-option="-P 3306"                      
    This option may be specified multiple times.
"""
    parser.add_option('--mysql-option', help=help,
                      action='append',  metavar='OPTION')

    help = """A command line option to pass to the mysqldump binary.
    Pass these unchanged, or quoted if you need to use a space,
e.g.:                                     
    --mysqldump-option=--ignore-table=foo                     
    --mysqldump-option=-Q                            
    --mysqldump-option="-P 3306"                      
    This option may be specified multiple times.
"""
    parser.add_option('--mysqldump-option', help=help,
                      action='append',  metavar='OPTION')

    return parser


def matches_mysql_pattern(pattern, string):
    """Check if a string matches a pattern containing MySQL wildcards.

    It converts a pattern containing MySQL wildcards into a normal regular
    expression, and matches the string to that:
    %  -> .*
    _  -> .
    \_ -> _

    Keyword arguments:
    pattern -- string to match which can contain MySQL wildcards (%, _)
    string -- string to try and match against the pattern

    returns -- True if string matches pattern, False otherwise
    """

    # replace MySQL % wildcard with regex .*
    pattern = pattern.replace('%', '.*')

    # temporarily replace escaped underscores with %
    pattern = pattern.replace('\_', '%')

    # replace unescaped underscored with regex .
    pattern = pattern.replace('_', '.')

    # undo temp replacement and convert % back into _
    pattern = pattern.replace('%', '_')
    if re.match(pattern, string):
        return True
    return False


def get_grant_info(grant_line, grant_re=[]):
    """Splits a grant line into a dictionary.

    This turns a grant line of the form:
    GRANT SELECT, INSERT ON *.* TO 'someuser'@'somehost';
    into a dictionary containing the fields:
    privileges - contains everything between 'GRANT' and 'ON
    database - name of a database or a wildcard (*)
    table - name of a database or a wildcard (*)
    user - username
    host - hostname
    grant_line - the original grant line used to build this dictionary

    Keyword arguments:
    grant_line -- string containg MySQL grant syntax
    grant_re -- used to hold static copy of grant regexp, should never be used

    returns -- dictionary containing parts of the grant line
    """

    # slightly hacky way of only compiling this regexp once
    if not grant_re:
        privs = '(.*)'
        db = '`?([^`]+)`?'
        table = '`?([^`]+)`?'
        user = "'?([^']*)'?"
        host = "'?([^']+)'?"

        grant = r'GRANT %s ON %s[.]%s TO %s@%s' \
                % (privs, db, table, user, host)
        grant_re.append(re.compile(grant))

    match = grant_re[0].match(grant_line)
    assert match, "GRANT regexp failed to match: %s" % grant_line

    info = {}
    info['privileges'] = match.group(1)
    info['database'] = match.group(2)
    info['table'] = match.group(3)
    info['user'] = match.group(4)
    info['host'] = match.group(5)
    info['grant_line'] = grant_line
    return info


def prompt_for_non_empty_value(prompt, error_msg):
    """Helper function which repeatedly prompts for a value until it receives
    a non empty one.
    
    Keyword arguments:
    prompt -- prompt to display to the user
    error_msg -- message displayed if an empty value is entered

    return -- string containing the value entered
    """

    value = raw_input(prompt)
    if value is None or value.strip() == '':
        print error_msg
        return prompt_for_non_empty_value(prompt, error_msg)
    return value

def prompt_for_value(prompt):
    """Helper function which prompts user for value based on input string

    The prompt should be of the form:
    Question? [a/B/c/d]:

    The values inside [] will be pulled out as accepted arguments. A
    capitalised value will be treated as the default, and returned if the
    user presses return at the prompt.

    Keyword arguments:
    prompt -- prompt to display to the user

    return -- string containing a single lowercase character
    """

    # pull out the text between [ and ]
    match = re.search('\[(.*)\]', prompt)
    assert match, 'Unable to pull out accepted args from prompt.'

    # split the text on /
    options = match.group(1).split('/')
    default_value = None
    accepted_options = []
    for option in options:
        # if the option is uppercase, use it as default value
        if option.isupper():
            option = option.lower()
            default_value = option
        accepted_options.append(option)

    # prompt user for input
    value = raw_input(prompt)

    # return the default value if the user pressed Return
    if value is None or value.strip() == '':
        return default_value

    # lowercase the value and return it is a valid choice
    lc_value = value[0].lower()
    if lc_value in accepted_options:
        return lc_value

    # otherwise warn the use and prompt for value again
    print "Invalid option '%s'." % value
    return prompt_for_value(prompt)


class MySQLHost(object):
    grant_re = None
    def __init__(self, host=None, user=None, password=None):
        """Connects to a database if all keyword arguments are given.

        Keyword arguments:
        host -- host to connect to MySQL server on
        user -- username to connect with
        password -- password to connect with
        """
        self.conn = None
        self.user = None
        self.password = None

        # empty passwords allowed
        if host and user and password is not None:
            self.connect(host, user, password)

    def connect(self, host, user, password):
        """Set up this object's connection to a MySQL server.

        Keyword arguments:
        host -- host to connect to MySQL server on
        user -- username to connect with
        password -- password to connect with
        """
        vprint("Connecting as %s to %s..." % (user, host), 1)
        self.conn = MySQLdb.connect(host=host, user=user, passwd=password)
        self.host = host
        self.user = user
        self.password = password
        vprint('done.\n', 1)

    def __del__(self):
        """Closes connection to MySQL server when done."""
        if self.conn:
            vprint("Closing DB connection to %s..." % self.host, 1)
            self.conn.close()
            vprint('done.\n', 1)

    def can_create_db(self, dbname):
        """Check whether the user we've connected with can create databases.

        Attempts to create and then drop a DB with a given name.  Will fail
        if unable to create a database, so should never drop an exisiting
        database.

        """

        assert dbname.strip() != ''

        # this requires writing to destination DB, so skip and assume
        # everythign is fine if --dry-run was specified
        if Settings.dry_run is True:
            return (True, None)

        # attempt create then drop DB - return False on any error
        c = self.conn.cursor()
        try:
            vprint('Checking we can create a DB on destination host...', 1)
            c.execute("CREATE DATABASE IF NOT EXISTS `%s`" % dbname)
            vprint('done.\n', 1)
            vprint('Checking we can drop the DB we created above on '
                   'destination host...', 1)
            c.execute("DROP DATABASE IF EXISTS `%s`" % dbname)
            vprint('done.\n', 1)
        except MySQLdb.Error, e:
            return (False, e)

        return (True, None)


    def get_user_grants(self, user, host):
        """Get a list of grant lines for a given user/host combination.

        The list returned contains the information returned by running:
        SHOW GRANTS FOR 'user'@'host';

        Keyword arguments:
        user -- MySQL username
        host -- hostname or MySQL wildcard (%)

        returns -- list of grant lines
        """
        c = self.conn.cursor()
        vprint("Getting grants for %s@%s ..." % (user, host), 1)
        c.execute('SHOW GRANTS FOR %s@%s', (user, host))
        vprint('done.\n', 1)
        rows = c.fetchall()
        c.close()
        return rows

    def get_all_grants(self):
        """Get list of all grants which exist on this MySQL server."""
        return self.get_db_grants(None)

    def get_db_grants(self, db_name):
        """Get a list of all grants which appear relevant to a database.

        Keyword arguments:
        db_name -- database name to match grants to

        returns -- list of MySQL grant lines
        """

        grant_lines = []
        user_rows = self.get_all_users()

        # get grants for all users
        grant_rows = []
        for user_row in user_rows:
            user = user_row[1]
            host = user_row[0]
            grant_rows.extend(self.get_user_grants(user, host))

        # if db_name is None, get all grants
        if db_name is None:
            for grant_row in grant_rows:
                grant_info = get_grant_info(grant_row[0])
                grant_lines.append(grant_info['grant_line'])
            return grant_lines

        # first get db specific grants...
        db_users = []
        for grant_row in grant_rows:
            grant_info = get_grant_info(grant_row[0])
            pattern = grant_info['database']
            if pattern != '*' and matches_mysql_pattern(pattern, db_name):
                grant_lines.append(grant_info['grant_line'])
                db_users.append(grant_info['user'])

        # ...then get wildcard grants
        for grant_row in grant_rows:
            grant_info = get_grant_info(grant_row[0])
            if grant_info['database'] == '*':
                if grant_info['user'] in db_users:
                    grant_lines.append(grant_info['grant_line'])

        return grant_lines

    def get_all_users(self):
        """Get a list all User,Host pairs from the mysql.user table"""
        c = self.conn.cursor()
        c.execute('SELECT Host, User FROM mysql.user')
        rows = c.fetchall()
        c.close()
        return rows

    def db_exists(self, database):
        """Check whether a database with a given name exists on this host.
        
        Keyword arguments:
        database -- case insensitive name of database to look for on this host

        returns -- True or False
        """

        found = False

        vprint("Checking if %s exists on %s..." % (database, self.host), 1)

        cursor = self.conn.cursor()
        cursor.execute('SHOW DATABASES LIKE %s;', database)
        if cursor.fetchone():
            found = True
        cursor.close()

        vprint('done.\n', 1)

        return found

    def dump_db_to_tempfile(self, database):
        """Use mysqldump to dump a database into a temporary file.
        
        Keyword arguments:
        database -- name of the database on this MySQL host to dump

        returns -- a File object
        """

        vprint('Dumping %s to temp file with mysqldump...' % database, 0)

        # flush output buffer to print the above message, dump may take time
        sys.stdout.flush()

        # Check if a different temp dir has been specified
        if Settings.temp_dir is None:
            dump_file = tempfile.NamedTemporaryFile()
        else:
            dump_file = tempfile.NamedTemporaryFile(dir=Settings.temp_dir)

        vprint("\nTemp file %s created.\n" % dump_file.name, 1)

        dump_args = Settings.mysqldump_extra_args + \
                    ["--host=%s" % self.host,
                     "--user=%s" % self.user,
                     '--password',
                     '--databases',
                     database,
                     ">%s" % dump_file.name]

        dump_cmd = '/bin/bash -c "%s %s"' % \
                   (Settings.mysqldump_path, ' '.join(dump_args))

        vprint("Executing mysqldump command:\n%s\n" % dump_cmd, 1)
        child = pexpect.spawn(dump_cmd, timeout=None)
        child.expect('Enter password:')
        child.sendline(self.password)
        vprint('done.\n', 1)

        # if mysqldump produces any output, it's an error, so save for later
        mysqldump_output = child.read()

        child.close()
        if child.exitstatus != 0:
            dump_file.close()
            sys.stderr.write("\nERROR: mysqldump exited with status %d:" % \
                             child.exitstatus)
            sys.stderr.write(mysqldump_output)
            sys.exit(1)

        vprint('done.', 0)
        return dump_file

    def execute_sql_from_file(self, file):
        """Execute SQL from a file using the MySQL binary.

        This uses the MySQL binary rather than python's MySQLdb cursor to get
        around the problem with exceeding the maximum query packet size the
        server accepts.

        Keyword arguments:
        file -- File object pointing to a MySQL dump file
        """

        # print first 15 lines of the file
        vprint("\nAbout to execute the following SQL on %s...\n" % self.host, 0)
        file.seek(0)
        max_lines = 15
        i = 0
        for line in file:
            if i < max_lines:
                vprint(line, 0)
                i += 1
            else:
                break

        # get the size of the file in human friendly format
        bytes = os.stat(file.name).st_size
        if bytes < 1024:
            file_size = "%f bytes" % bytes
        elif bytes < 1048576: # 1024^2
            file_size = "%.2f Kb" % (bytes/1024.0)
        else:
            file_size = "%.2f Mb" % (bytes/1048576.0)

        msg = """
----------------------------------------------------
The first 15 lines of the SQL to be executed are shown above. Please choose
from one of the following options:
f - display the full output of the dump (%s)
y - execute the above SQL on %s
n - cancel the database copy
""" % (file_size, self.host)
        vprint(msg, 0)

        # Check if we're in force mode
        if Settings.force is True:
            confirm = 'y'
        else:
            prompt = "Execute the above SQL on %s? [f/Y/n]: " % self.host
            confirm = prompt_for_value(prompt)
            
        if confirm == 'f':
            assert Settings.force is False, 'confirm should not = f if force'
            file.seek(0)
            print file.read()
            prompt = "Execute the above SQL on %s? [Y/n]: " % self.host
            confirm = prompt_for_value(prompt)

        if confirm == 'y':
            mysql_args = Settings.mysql_extra_args + \
                         ["--host=%s" % self.host,
                          "--user=%s" % self.user,
                          '--password',
                          "<%s" % file.name]
            mysql_cmd = '/bin/bash -c "%s %s"' % \
                        (Settings.mysql_path, ' '.join(mysql_args))
            vprint("Executing mysql command:\n%s\n" % mysql_cmd, 1)

            # do nothing if it's a dry run
            if Settings.dry_run is False:
                child = pexpect.spawn(mysql_cmd, timeout=None)
                child.expect('Enter password:')
                child.sendline(self.password)

                # if mysql any output, it's an error, so save for later
                mysql_output = child.read()

                child.close()
                if child.exitstatus != 0:
                    msg = "\nERROR: mysql exited with status %d:" % \
                          child.exitstatus
                    sys.stderr.write(msg)
                    sys.stderr.write(mysql_output)
                    sys.exit(1)

            vprint('done.\n', 1)

            # if dry run, or mysql exited with status 0
            vprint('done.\n', 0)
            return

        # if user entered something other than 'y'
        assert confirm == 'n', 'confirm should equal y/n'
        vprint('Database copy cancelled.\n', 0)
        sys.exit(0)

    def copy_db_grants_to(self, database, dest):
        """Copy grants from a database on this host to another host.

        This searches through all user grants, and tries to pick the most
        relevant ones to execute on the destination server.  User will be
        prompted before each grant is executed.

        Keyword arguments:
        database -- name of the database on this host
        dest -- MySQLHost object which is connected to another MySQL server
        """

        vprint("\nGetting relevant privileges for DB %s on %s... " % \
              (database, self.host), 0)

        grants = self.get_db_grants(database)
        all_grants = False
        if not grants:
            vprint('\nWARNING: No DB specific grants found.  This is '
                   'likely to mean that this database is accessed through '
                   'an account with global privileges.', 0)

            if settings.force is False:
                p = 'Get ALL grants and prompt for copy of each? [y/N]: '
                do_all = prompt_for_value(p)
                if do_all == 'y':
                    grants = self.get_all_grants()
                    all_grants = True

        if grants:
            grants.reverse()
        vprint('done.\n', 0)

        n_grants = len(grants)
        msg = """
Found %d possibly relevant grants...
Please choose one of the following options for each grant:
y - apply the grant on %s
n - do not apply this grant on %s
l - list all grants that appear relevant
q - quit, do not apply any more grants
""" % (n_grants, dest.host, dest.host)
        vprint(msg, 0)

        cursor = dest.conn.cursor()
        done = False
        grant_num = 0
        for grant in grants:
            grant_num += 1
            if done:
                break

            while True:
                gtitle = "----- Grant %d of %d -----\n" % (grant_num, n_grants)
                vprint(gtitle, 0)
                vprint(grant + '\n', 0)

                if Settings.force is True:
                    # don't copy all grants in force mode
                    if all_grants:
                        confirm = 'n'
                    else:
                        confirm = 'y'
                else:
                    prompt = "Execute the above grant on %s? [Y/n/l/q]: " % \
                             dest.host
                    confirm = prompt_for_value(prompt)

                if confirm == 'y':
                    if Settings.dry_run is False:
                        cursor.execute(grant + ';')
                    break
                elif confirm == 'n':
                    break
                elif confirm == 'q':
                    done = True
                    break
                else:
                    assert confirm == 'l', 'confirm should be one of [y/n/l/q]'
                    print '----------'
                    print 'All possibly relevant grants:'
                    print '\n'.join(grants)

        if Settings.dry_run is False:
            cursor.execute('FLUSH PRIVILEGES;')

        cursor.close()
        vprint('done.', 0)

def prompt_for_fields(prefix, options, fields):
    """Helper function to prompt for fields which weren't specified on command
    line
    """
    values = {}

    for field in fields:
        prompt = "%s %s: " % (prefix, field)
        optval = eval("options.%s_%s" % (prefix, field))
        if optval:
            values[field] = optval
            if field is 'password':
                optval = '<hidden>'
            vprint("%s%s\n" % (prompt, optval), 0)
        else:
            if field is 'password':
                values[field] = getpass.getpass(prompt)
            else:
                error_msg = "%s cannot be empty, please enter a new value." \
                            % field
                values[field] = prompt_for_non_empty_value(prompt, error_msg)
    return values


def main():
    # parse command line args
    parser = get_option_parser()
    (options, args) = parser.parse_args()

    opt_dict = vars(options)
    options_set = [key for key in opt_dict if opt_dict[key] is not None]

    # print brief help message if no options given
    if options_set == []:
        parser.print_help()
        return 0

    # modify global settings from command line args if needed
    if options.force:
        Settings.force = True

    if options.quiet:
        Settings.verbosity = -1

    if options.verbose: 
        Settings.verbosity = 1

    if options.dry_run:
        Settings.dry_run = True

    if options.temp_dir:
        Settings.temp_dir = options.temp_dir

    if options.mysql_option:
        Settings.mysql_extra_args = options.mysql_option

    if options.mysqldump_option:
        Settings.mysqldump_extra_args = options.mysqldump_option

    # check for existence of mysql and mysqldump
    if not os.access(Settings.mysql_path, os.F_OK|os.R_OK|os.X_OK):
        sys.stderr.write("ERROR: Unable execute mysql binary at '%s'\n" %\
                         Settings.mysql_path)
        sys.exit(1)

    if not os.access(Settings.mysqldump_path, os.F_OK|os.R_OK|os.X_OK):
        sys.stderr.write("ERROR: Unable execute mysqldump binary at '%s'\n" %\
                         Settings.mysqldump_path)
        sys.exit(1)

    # check that user specified tmp dir is readable/writable
    if Settings.temp_dir is not None:
        if not os.access(Settings.temp_dir, os.F_OK|os.R_OK|os.W_OK):
            sys.stderr.write(
                "Error: unable to find writable temp dir at '%s'\n" \
                % Settings.temp_dir)
            sys.exit(1)
        

    # prompt for source information missing from options
    source_info = prompt_for_fields('source', options,
                               ['host', 'user', 'password', 'database'])

    # prompt for destination information missing from options
    dest_info = prompt_for_fields('destination', options,
                             ['host', 'user', 'password'])

    try:
        # set up DB connections to both source and destination DB
        source_host = MySQLHost(source_info['host'], source_info['user'],
                                source_info['password'])
        dest_host = MySQLHost(dest_info['host'], dest_info['user'],
                              dest_info['password'])

        # Check that DB to copy doesn't already exist on destination host
        if dest_host.db_exists(source_info['database']):
            msg = """
ERROR: a database named '%s' already exists on %s.
Please rename or drop it before attempting to copy another database with the
same name.
""" % (source_info['database'], dest_host.host)
            sys.stderr.write(msg)
            sys.exit(1)

        # check that we're able to create DB on destination host
        (can_create, error) = dest_host.can_create_db(source_info['database'])
        if can_create is False:
            msg = """
Unable to create database on '%s'.  Please check that the user you have
connected with has sufficient privileges to create databases.
Error message was: %s
""" % (dest_host.host, error)
            sys.stderr.write(msg)
            sys.exit(1)
    except MySQLdb.Error, e:
        sys.stderr.write("MySQL error: %s\n" % e)
        sys.exit(1)

    # use mysqldump to write temporary dumpfile
    dump_file = source_host.dump_db_to_tempfile(source_info['database'])
    assert dump_file, 'should have checked for mysqldump failure earlier'

    dest_host.execute_sql_from_file(dump_file)

    try:
        source_host.copy_db_grants_to(source_info['database'], dest_host)
    except MySQLdb.Error, e:
        sys.stderr.write("MySQL error: %s\n" % e)
        sys.exit(1)

    if Settings.dry_run:
        vprint('\nDry run finished, no data modified on hosts.', 0)
    else:
        vprint('\nAll done, database and selected privileges copied.', 0)
    sys.exit(0)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        vprint('\nCancelled by keyboard interrupt.\n', 0)
        sys.exit(0)
